#lang racket

(require "util/utils.scm")
(require "util/gensym.scm")
(provide pass-denest)

(define (pass-denest prg)
  (let ((definitions (cdr (assoc 'definitions prg)))
        (code (cdr (assoc 'code prg))))
    `((definitions . ,definitions)
      (procedures . ,(map* pass-denest-code code)))))

(define (pass-denest-code code)
  (let ((label (car code))
        (env (cadr code))
        (binds (caddr code))
        (exp (cadddr code)))
    `(,label ,env ,binds ,(make-let* (pass-denest-exp exp #f)))))

;; TODO add a parameter to say we don't need the return value
;;      so that we can avoid creating temporaries for builtin
(define (pass-denest-exp exp simple)
  (match exp
    (`(atomic ,exp) (cons '() `(atomic ,exp)))
    (`(var ,type ,x) (cons '() `(var ,type ,x)))

    ;; copy paste criminal
    (`(app . ,exps)
     (let* ((work-result (mapply pass-denest-exp exps #t))
	    (work (concatenate (map car work-result)))
	    (results (map cdr work-result))
	    (tmp (dgensym 't)))
       (if simple
	   (cons (append work (list `(,tmp (app . ,results))))
		 `(var tmp ,tmp))
	   (cons work `(app . ,results)))))
    ;; forgive me..
    (`(known-app ,i . ,exps)
     (let* ((work-result (mapply pass-denest-exp exps #t))
	    (work (concatenate (map car work-result)))
	    (results (map cdr work-result))
	    (tmp (dgensym 't)))
       (if simple
	   (cons (append work (list `(,tmp (known-app ,i . ,results))))
		 `(var tmp ,tmp))
	   (cons work `(known-app ,i . ,results)))))
    
    (`(begin . ,exps)
     (let loop ((exps exps) (binds '()) (tmp #f))
       (if (null? exps)
           (if tmp
               (cons binds tmp)
               (error 'pass-denest-exp "empty begin"))
           (let ((w-r (pass-denest-exp (car exps)
                                       (or simple (not (null? (cdr exps)))) ;; only the last expression in a BEGIN can end in non-simple work
                                       )))                                  ;; plus, if a simple expression is required we have no choice
             (loop (cdr exps) (append binds (car w-r)) (cdr w-r))))))
    (`(closure ,env ,label)
     (let ((tmp (dgensym 't)))
       (cons `((,tmp (closure ,env ,label)))
             `(var tmp ,tmp))))
    (`(quote ,x) (cons '() `(quote ,x)))
    (`(if ,t ,c ,a)
      (let* ((t-wr (pass-denest-exp t #t)) ;; this must be simple otherwise we may have (if (if ...) ...)
             (t-work (car t-wr))
             (t-result (cdr t-wr))
             (c-let (make-let* (pass-denest-exp c #f))) ;; these need not be simple because we're turning them into their own
             (a-let (make-let* (pass-denest-exp a #f))) ;; let expressions
             (tmp (dgensym 't)))
        (if simple
            (cons (append t-work (list (list tmp `(if ,t-result ,c-let ,a-let)))) `(var tmp ,tmp))
            (cons t-work `(if ,t-result ,c-let ,a-let)))))
    (`(let* ,binds ,body)
     (let loop ((binds binds) (work '()))
       (if (null? binds)
           (let ((w-r (pass-denest-exp body simple)))
             (cons (append work (car w-r)) (cdr w-r)))
           (let* ((tmp (caar binds))
                  (exp (cadar binds))
                  (w-r (pass-denest-exp exp #t)))
             (loop (cdr binds) (append work (car w-r) (list (list tmp (cdr w-r)))))))))
    (`(builtin ,nm . ,args)
     (let* ((work-result (mapply pass-denest-exp args #t))
            (work (concatenate (map car work-result)))
            (results (map cdr work-result))
            (tmp (dgensym 't)))
       (if simple
           (cons (append work (list (list tmp `(builtin ,nm . ,results)))) `(var tmp ,tmp))
           (cons work `(builtin ,nm . ,results)))))))

(define (make-let* work-result)
  (if (null? (car work-result))
      (cdr work-result)
      `(let* ,(car work-result) ,(cdr work-result))))
